import os

import numpy as np
import torch
import cv2
from torch import load, argmax, reshape
from torch import nn
from torchvision import transforms
from torch.utils.tensorboard import SummaryWriter
from matplotlib import pyplot as plt

# %%
class AlexNet(nn.Module):
    def __init__(self):
        super(AlexNet, self).__init__()
        self.conv = nn.Sequential(
            # in_channels, out_channels, kernel_size, stride
            nn.Conv2d(3, 96, 11, 4),
            nn.ReLU(),
            # kernel_size, stride
            nn.MaxPool2d(3, 2),

            # 减小卷积窗口，使用步长为2来使得输入与输出的高和宽一致，
            # 且增大输出通道数
            nn.Conv2d(96, 256, 5, 1, 2),
            nn.ReLU(),
            nn.MaxPool2d(3, 2),
            # 连续3个卷积层，且使用更小的卷积窗口。除了最后的卷积层外，
            # 进一步增大了输出通道数。
            # 前两个卷积层后不使用池化层来减小输入的高和宽
            nn.Conv2d(256, 384, 3, 1, 1),
            nn.ReLU(),
            nn.Conv2d(384, 384, 3, 1, 1),
            nn.ReLU(),
            nn.Conv2d(384, 256, 3, 1, 1),
            nn.ReLU(),
            nn.MaxPool2d(3, 2)
        )
        # 这里全连接层的输出个数比LeNet中的大数倍。使用丢弃层来缓解过拟合
        self.fc = nn.Sequential(
            nn.Linear(256 * 5 * 5, 4096),
            nn.ReLU(),
            nn.Dropout(0.5),
            nn.Linear(4096, 4096),
            nn.ReLU(),
            nn.Dropout(0.5),
            # 输出层。由于这里使用Fashion-MNIST
            # 所以用类别数为10，而非论文中的1000
            nn.Linear(4096, 2)
        )

    def forward(self, img):
        feature = self.conv(img)
        output = self.fc(feature.view(img.shape[0], -1))
        return output


net = AlexNet()
net.load_state_dict(torch.load('E:\Class_exprience\Machine_Vision\9_class_design\model16.pt'), True)

# checkpoint = torch.load(modelpath)  #modelpath是你要加载训练好的模型文件地址
# model.load_state_dict(checkpoint['state_dict'])
# output = model(x)

# %%测试

def classfiey(path):
    # path=r'E:\Class_exprience\Machine_Vision\9_class_design\training_set\training_set\cats\cat.2.jpg'
    # path=r"E:\Class_exprience\Machine_Vision\9_class_design\test_set\test_set\dogs\dog.4001.jpg"
    resize = 224
    path = os.path.abspath(path)
    transform = transforms.Compose([
        transforms.ToPILImage(),
        transforms.Resize((resize, resize)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ])

    img = cv2.imread(path)
    img_1 = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    img_2 = transform(img_1)
    image = reshape(img_2, (1, 3, resize, resize))

    net.eval()
    result = net(image)
    selection = ['猫', '狗']
    return selection[argmax(result)]


# writer=SummaryWriter('./log')
# input = torch.rand(13, 3, 100, 100)
# writer.add_graph(net,input)
# writer.close()
# plt.figure(1)
# plt.imshow(img)
# plt.show()
# plt.figure(2)
# plt.imshow(img_1)
# plt.show()
# plt.figure(3)
# plt.imshow(img_2)
# plt.show()
# plt.figure(4)
# plt.imshow(image)
# plt.show()




